package nsfs

import (
	"GLibHac/crypto"
	"GLibHac/fssystem"
	"GLibHac/tools"
	"encoding/binary"
	"fmt"
	"github.com/pkg/errors"
)

type NcaSectionType byte

func (n NcaSectionType) String() string {
	names := []string{
		"Code",
		"Data",
		"Logo",
	}
	return names[n]
}

type NcaKeyType byte

const (
	NcaKeyTypeAesXts0 = NcaKeyType(iota)
	NcaKeyTypeAesXts1
	NcaKeyTypeAesCtr
	NcaKeyTypeAesCtrEx
	NcaKeyTypeAesCtrHw
)

const (
	NcaSectionTypeCode = NcaSectionType(iota)
	NcaSectionTypeData
	NcaSectionTypeLogo
)

type IntegrityCheckLevel byte

func (b IntegrityCheckLevel) String() string {
	names := []string{
		"None",
		"IgnoreOnInvalid",
		"ErrorOnInvalid",
	}
	return names[b]
}

const (
	IntegrityCheckLevelNone            = IntegrityCheckLevel(iota) // No integrity checks will be performed.
	IntegrityCheckLevelIgnoreOnInvalid                             // Invalid blocks will be marked as invalid when read, and will not cause an error.
	IntegrityCheckLevelErrorOnInvalid                              //An error will be thrown if an integrity check fails.
)

var KakNames = []string{"application", "ocean", "system"}

type IntegrityStorageType byte

func (i IntegrityStorageType) String() string {
	names := []string{
		"Save",
		"RomFs",
		"PartitionFs",
	}
	return names[i]
}

const (
	IntegrityStorageTypeSave = IntegrityStorageType(iota)
	IntegrityStorageTypeRomFs
	IntegrityStorageTypePartitionFs
)

type IntegrityVerificationInfo struct {
	Data      fssystem.IStorage
	BlockSize int32
	IntegrityStorageType
}

func NewIntegrityVerificationInfo(storage fssystem.IStorage, blockSize int32, integrityStorageType IntegrityStorageType) *IntegrityVerificationInfo {
	return &IntegrityVerificationInfo{
		Data:                 storage,
		BlockSize:            blockSize,
		IntegrityStorageType: integrityStorageType,
	}
}

type Nca struct {
	Header              *NcaHeader
	nca0TransformedBody fssystem.IStorage
	baseStorage         fssystem.IStorage
	keySet              *KeySet
	Nca0KeyArea         []byte
}

func (n *Nca) SectionExistsByIndex(index int32) bool {
	return n.Header.IsSectionEnabled(index)
}

func (n *Nca) SectionExists(sectionType NcaSectionType) bool {
	index := TryGetSectionIndexFromType(sectionType, n.Header.ContentType)
	if index == -1 {
		return false
	}
	return n.Header.IsSectionEnabled(index)
}

func (n *Nca) OpenFileSystemByType(sectionType NcaSectionType, integrityCheckLevel IntegrityCheckLevel) (IFileSystem, error) {
	index, err := n.GetSectionIndexFromType(sectionType)
	if err != nil {
		return nil, err
	}
	return n.OpenFileSystem(index, integrityCheckLevel)
}

func (n *Nca) OpenFileSystem(index int32, integrityCheckLevel IntegrityCheckLevel) (IFileSystem, error) {
	storage, err := n.OpenStorage(index, integrityCheckLevel, false)
	if err != nil {
		return nil, err
	}
	fsHeader, err := n.GetFsHeader(index)
	if err != nil {
		return nil, err
	}
	return OpenFileSystem(storage, fsHeader)
}

func OpenFileSystem(storage fssystem.IStorage, header *NcaFsHeader) (IFileSystem, error) {
	switch header.FormatType {
	case NcaFormatTypePfs0:
		return NewPartitionFileSystem(storage)
	case NcaFormatTypeRomfs:
		return NewRomFsFileSystem(storage)
	default:
		return nil, errors.Errorf("unknown format type: %d", header.FormatType)
	}
}

func (n *Nca) OpenStorage(index int32, integrityCheckLevel IntegrityCheckLevel, leaveCompressed bool) (fssystem.IStorage, error) {
	rawStorage, err := n.OpenRawStorage(index, false)
	if err != nil {
		return nil, err
	}
	header, err := n.Header.GetFsHeader(index)
	if err != nil {
		return nil, err
	}
	if header.EncryptionType == NcaEncryptionTypeAesCtrEx {
		return rawStorage.SliceRange(0, header.PatchInfo.RelocationTreeOffset), nil
	}

	returnStorage, err := n.CreateVerificationStorage(integrityCheckLevel, header, rawStorage)

	if !leaveCompressed && header.ExistsCompressionLayer() {
		returnStorage, err = n.OpenCompressedStorage(header, returnStorage)
	}

	return returnStorage, err
}

func (n *Nca) GetSectionIndexFromType(sectionType NcaSectionType) (int32, error) {
	return GetSectionIndexFromType(sectionType, n.Header.ContentType)
}

func (n *Nca) GetFsHeader(index int32) (*NcaFsHeader, error) {
	if n.Header.IsNca0() {
		return n.GetNca0FsHeader(index)
	}
	return n.Header.GetFsHeader(index)
}

func (n *Nca) OpenRawStorage(index int32, openEncrypted bool) (fssystem.IStorage, error) {
	if n.Header.IsNca0() {
		return n.OpenNca0RawStorage(index, openEncrypted)
	}

	storage, err := n.OpenSectionStorage(index)
	if err != nil {
		return nil, err
	}

	if n.Header.IsEncrypted == openEncrypted {
		return storage, nil
	}

	return n.OpenDecryptedStorage(storage, index, !openEncrypted)
}

func (n *Nca) OpenNca0RawStorage(index int32, openEncrypted bool) (fssystem.IStorage, error) {
	if !n.SectionExistsByIndex(index) {
		return nil, errors.Errorf("cannot open Nca0 storage for index [ %d ], section not exist", index)
	}
	offset := n.Header.GetSectionStartOffset(index) - 0x400
	size := n.Header.GetSectionSize(index)

	bodyStorage, err := n.OpenNca0BodyStorage(openEncrypted)
	if err != nil {
		return nil, err
	}
	baseSize := bodyStorage.Size()

	if !IsSubRange(offset, size, baseSize) {
		return nil, errors.Errorf("section offset (0x%x) and length (0x%x) fall outside the total NCA length (0x%x)",
			offset+0x400,
			size,
			baseSize+0x400)
	}
	return bodyStorage.SliceRange(offset, size), nil
}

func IsSubRange(startIndex int64, subLength int64, length int64) bool {
	isOutOfRange := startIndex < 0 || startIndex > length || subLength < 0 || startIndex > length-subLength
	return !isOutOfRange
}

func (n *Nca) OpenSectionStorage(index int32) (fssystem.IStorage, error) {
	if !n.SectionExistsByIndex(index) {
		return nil, errors.Errorf("section does not exist, index=%d", index)
	}

	offset := n.Header.GetSectionStartOffset(index)
	size := n.Header.GetSectionSize(index)

	ncaStorageSize := n.baseStorage.Size()

	fsHeader, err := n.Header.GetFsHeader(index)
	if err != nil {
		return nil, err
	}

	if fsHeader.ExistsSparseLayer() {
		sparseInfo := fsHeader.NcaSparseInfo
		header := NewBucketTreeHeader(sparseInfo.MetaHeader)
		if err = header.Verify(); err != nil {
			return nil, err
		}
		//TODO need to implement
		return nil, errors.Errorf("not implemented")
	}
	if !n.IsSubRange(offset, size, ncaStorageSize) {
		return nil, errors.Errorf("section offset (0x%x) and length (0x%x) fall outside the total NCA length (0x%x)",
			offset,
			size,
			ncaStorageSize)
	}
	return n.baseStorage.SliceRange(offset, size), nil
}

func (n *Nca) OpenDecryptedStorage(storage fssystem.IStorage, index int32, decrypting bool) (fssystem.IStorage, error) {
	header, err := n.GetFsHeader(index)
	if err != nil {
		return nil, err
	}
	switch header.EncryptionType {
	case NcaEncryptionTypeNone:
		return storage, nil
	case NcaEncryptionTypeXTS:
		return n.OpenAesXtsStorage(storage, index, decrypting)
	case NcaEncryptionTypeAesCtr:
		offset := n.Header.GetSectionStartOffset(index)
		return n.OpenAesCtrStorage(storage, index, offset, header.UpperCounter)
	case NcaEncryptionTypeAesCtrEx:
		return n.OpenAesCtrExStorage(storage, index, decrypting)
	default:
		return nil, errors.Errorf("unknown encryption type: %d", header.EncryptionType)
	}
}

func (n *Nca) OpenNca0BodyStorage(openEncrypted bool) (fssystem.IStorage, error) {
	// NCA0 encrypts the entire NCA body using AES-XTS instead of
	// using different encryption types and IVs for each section.

	if openEncrypted == n.Header.IsEncrypted {
		return n.GetRawStorage(), nil
	}

	if n.nca0TransformedBody != nil {
		return n.nca0TransformedBody, nil
	}

	key0, err := n.GetContentKey(NcaKeyTypeAesXts0)
	if err != nil {
		return nil, err
	}
	key1, err := n.GetContentKey(NcaKeyTypeAesXts1)
	if err != nil {
		return nil, err
	}
	rawStorage := n.GetRawStorage()
	nca0TransformedBody, err := NewAes128XtsStorage(rawStorage, key0, key1, NcaHeaderSectorSize, !openEncrypted)
	if err != nil {
		return nil, err
	}
	n.nca0TransformedBody = nca0TransformedBody
	return nca0TransformedBody, nil
}

func (n *Nca) GetRawStorage() fssystem.IStorage {
	ncaSize := n.baseStorage.Size()
	return n.baseStorage.SliceRange(0x400, ncaSize-0x400)
}

func (n *Nca) GetContentKey(ncaKeyType NcaKeyType) ([]byte, error) {
	if n.Header.HasRightsId() {
		return n.GetDecryptedTitleKey()
	} else {
		return n.GetDecryptedKey(int(ncaKeyType))
	}
}

func (n *Nca) GetDecryptedTitleKey() ([]byte, error) {
	keyRevision := tools.GetMasterKeyRevision(n.Header.GetKeyGeneration())
	titleKek := n.keySet.TitleKeks[keyRevision].Data

	var rightsId = string(n.Header.RightsId)
	accessKey := n.keySet.ExternalKeySet[rightsId]
	if accessKey == nil {
		return nil, errors.Errorf("missing NCA title key: %s", rightsId)
	}

	if tools.IsZeroArray(titleKek) {
		return nil, errors.Errorf("unable to decrypt title key: titlekek_%02x", keyRevision)
	}

	encryptedKey := accessKey.Data
	return crypto.AesECBDecrypt(encryptedKey, titleKek)
}

func (n *Nca) GetDecryptedKey(index int) ([]byte, error) {
	if index < 0 || index > 3 {
		return nil, errors.Errorf("key index [%d] out of range, must be between 0 and 3", index)
	}

	// Handle old NCA0s that use different key area encryption
	if n.Header.FormatVersion == NcaVersionNca0FixedKey || n.Header.FormatVersion == NcaVersionNca0RsaOaep {
		data, err := n.GetDecryptedKeyAreaNca0()
		if err != nil {
			return nil, err
		}
		offset := 0x10 * index
		return data[offset : offset+0x10], nil
	}

	keyRevision := tools.GetMasterKeyRevision(n.Header.GetKeyGeneration())
	keyIndex := n.Header.KeyAreaKeyIndex
	keyAreaKey := n.keySet.KeyAreaKeys[keyRevision][keyIndex].Data

	if tools.IsZeroArray(keyAreaKey) {
		return nil, errors.WithStack(fmt.Errorf("unable to decrypt NCA section: key_area_key_%s_%02x", KakNames[keyIndex], keyRevision))
	}

	encryptedKey, err := n.Header.GetEncryptedKey(index)
	if err != nil {
		return nil, err
	}
	return crypto.AesECBDecrypt(encryptedKey, keyAreaKey)
}

func (n *Nca) GetDecryptedKeyAreaNca0() ([]byte, error) {
	if n.Nca0KeyArea != nil {
		return n.Nca0KeyArea, nil
	}

	if n.Header.FormatVersion == NcaVersionNca0FixedKey {
		n.Nca0KeyArea = n.Header.GetKeyArea()
	} else {
		return nil, errors.Errorf("unsupported NCA version: %d", n.Header.FormatVersion)
	}

	return n.Nca0KeyArea, nil
}

func (n *Nca) CreateVerificationStorage(integrityCheckLevel IntegrityCheckLevel, header *NcaFsHeader, rawStorage fssystem.IStorage) (fssystem.IStorage, error) {
	switch header.HashType {
	case NcaHashTypeSha256:
		return n.InitIvfcForPartitionFs(header.NcaFsIntegrityInfoSha256, rawStorage, integrityCheckLevel)
	case NcaHashTypeIvfc:
		// The FS header of an NCA0 section with IVFC verification must be manually skipped
		if n.Header.IsNca0() {
			rawStorage = rawStorage.Slice(0x200)
		}
		return n.InitIvfcForRomFs(header.NcaFsIntegrityInfoIvfc, rawStorage, integrityCheckLevel)
	default:
		return nil, errors.Errorf("invalid raw storage hash type: : %s", header.HashType)
	}
}

func (n *Nca) OpenCompressedStorage(header *NcaFsHeader, baseStorage fssystem.IStorage) (fssystem.IStorage, error) {
	compressionInfo := header.CompressionInfo
	bucketTreeHeader := compressionInfo.MetaHeader
	if err := bucketTreeHeader.Verify(); err != nil {
		return nil, err
	}
	nodeStorageSize, err := QueryNodeStorageSize(bucketTreeHeader.EntryCount)
	if err != nil {
		return nil, err
	}
	var entryStorageSize = QueryEntryStorageSize(bucketTreeHeader.EntryCount)
	var tableOffset = compressionInfo.MetaOffset
	var tableSize = compressionInfo.MetaSize

	if entryStorageSize+nodeStorageSize > tableSize {
		return nil, errors.Errorf("entry size %d exceeds table size limit %d", entryStorageSize+nodeStorageSize, tableSize)
	}
	var dataStorage = baseStorage.SliceRange(0, tableOffset)
	var nodeStorage = baseStorage.SliceRange(tableOffset, nodeStorageSize)
	var entryStorage = baseStorage.SliceRange(tableOffset+nodeStorageSize, entryStorageSize)
	compressedStorage, err := NewCompressedStorage(dataStorage, nodeStorage, entryStorage, bucketTreeHeader.EntryCount)
	if err != nil {
		return nil, err
	}
	return NewCachedStorage(compressedStorage, 0x4000, 32), nil
}

func (n *Nca) InitIvfcForPartitionFs(info *NcaFsIntegrityInfoSha256, pfsStorage fssystem.IStorage, integrityCheckLevel IntegrityCheckLevel) (fssystem.IStorage, error) {
	offset, err := info.GetLevelOffset(0)
	if err != nil {
		return nil, err
	}
	size, err := info.GetLevelSize(0)
	if err != nil {
		return nil, err
	}
	hashStorage := pfsStorage.SliceRange(offset, size)
	offset, err = info.GetLevelOffset(1)
	if err != nil {
		return nil, err
	}
	size, err = info.GetLevelSize(1)
	if err != nil {
		return nil, err
	}
	dataStorage := pfsStorage.SliceRange(offset, size)

	var initInfo = make([]*IntegrityVerificationInfo, 3)

	// Set the master hash
	initInfo[0] = NewIntegrityVerificationInfo(
		//TODO Get hash directly from header
		fssystem.NewMemoryStorage(info.GetMasterHash()),
		0,
		IntegrityStorageTypePartitionFs,
	)
	size, err = info.GetLevelSize(0)
	if err != nil {
		return nil, err
	}
	initInfo[1] = NewIntegrityVerificationInfo(
		hashStorage,
		int32(size),
		IntegrityStorageTypePartitionFs,
	)
	initInfo[2] = NewIntegrityVerificationInfo(dataStorage,
		info.BlockSize,
		IntegrityStorageTypePartitionFs,
	)
	return NewHierarchicalIntegrityVerificationStorage(initInfo, integrityCheckLevel)
}

func (n *Nca) InitIvfcForRomFs(ivfc *NcaFsIntegrityInfoIvfc, dataStorage fssystem.IStorage, integrityCheckLevel IntegrityCheckLevel) (fssystem.IStorage, error) {
	var initInfo = make([]*IntegrityVerificationInfo, ivfc.LevelCount)

	initInfo[0] = NewIntegrityVerificationInfo(
		fssystem.NewMemoryStorage(ivfc.GetMasterHash()),
		0,
		IntegrityStorageTypeSave)

	for i := int32(1); i < ivfc.LevelCount; i++ {
		levelInfo, err := ivfc.GetLevelInfo(i - 1)
		if err != nil {
			return nil, err
		}
		initInfo[i] = NewIntegrityVerificationInfo(
			dataStorage.SliceRange(levelInfo.Offset, levelInfo.Size),
			1<<levelInfo.BlockSize,
			IntegrityStorageTypeRomFs,
		)
	}
	return NewHierarchicalIntegrityVerificationStorage(initInfo, integrityCheckLevel)
}

func (n *Nca) GetNca0FsHeader(index int32) (*NcaFsHeader, error) {
	// NCA0 stores the FS header in the first block of the section instead of the header
	bodyStorage, err := n.OpenNca0BodyStorage(false)
	if err != nil {
		return nil, err
	}
	offset := n.Header.GetSectionStartOffset(index) - 0x400
	fsHeaderData := make([]byte, 0x200)
	if _, err := bodyStorage.Read(offset, fsHeaderData); err != nil {
		return nil, err
	}
	return NewNcaFsHeader(fsHeaderData), nil
}

func (n *Nca) OpenAesXtsStorage(baseStorage fssystem.IStorage, index int32, decrypting bool) (*Aes128XtsStorage, error) {
	sectorSize := int32(0x200)

	key0, err := n.GetContentKey(NcaKeyTypeAesXts0)
	if err != nil {
		return nil, err
	}
	key1, err := n.GetContentKey(NcaKeyTypeAesXts1)
	if err != nil {
		return nil, err
	}
	// todo: Handle xts for nca version 3
	return NewAes128XtsStorage(baseStorage, key0, key1, sectorSize, decrypting)
}

func (n *Nca) OpenAesCtrStorage(baseStorage fssystem.IStorage, index int32, offset int64, upperCounter uint64) (*Aes128CtrStorage, error) {
	key, err := n.GetContentKey(NcaKeyTypeAesCtr)
	if err != nil {
		return nil, err
	}
	counter := CreateAes128CtrCounter(upperCounter, n.Header.GetSectionStartOffset(index))
	return NewAes128CtrStorage(baseStorage, key, offset, counter)
}

func (n *Nca) OpenAesCtrExStorage(baseStorage fssystem.IStorage, index int32, decrypting bool) (fssystem.IStorage, error) {
	//TODO: need to implement
	panic("implement me")
}

func (n *Nca) IsSubRange(startIndex int64, subLength int64, length int64) bool {
	isOutOfRange := startIndex < 0 || startIndex > length || subLength < 0 || startIndex > length-subLength
	return !isOutOfRange
}

func CreateAes128CtrCounter(counter uint64, offset int64) []byte {
	counterBytes := make([]byte, 0x10)
	binary.BigEndian.PutUint64(counterBytes, counter)
	binary.BigEndian.PutUint64(counterBytes[8:], uint64(offset/0x10))
	return counterBytes
}

func NewNca(keySet *KeySet, reader fssystem.IStorage) (nca *Nca, err error) {
	nca = &Nca{
		baseStorage: reader,
		keySet:      keySet,
	}
	nca.Header, err = NewNcaHeader(keySet, reader)
	return
}

func GetSectionIndexFromType(sectionType NcaSectionType, contentType NcaContentType) (int32, error) {
	index := TryGetSectionIndexFromType(sectionType, contentType)
	if index < 0 {
		return -1, errors.Errorf("NCA does not contain this section type: {type= %s, contentType=%s}", sectionType, contentType)
	}
	return index, nil
}

func TryGetSectionIndexFromType(sectionType NcaSectionType, contentType NcaContentType) int32 {
	if sectionType == NcaSectionTypeCode && contentType == NcaContentTypeProgram {
		return 0
	}
	if sectionType == NcaSectionTypeData && contentType == NcaContentTypeProgram {
		return 1
	}

	if sectionType == NcaSectionTypeLogo && contentType == NcaContentTypeProgram {
		return 2
	}
	if sectionType == NcaSectionTypeData {
		return 0
	}
	return -1
}
